#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;
template<typename T_VAR, typename T_INDICES, typename T_UPDATES> class KernelScatterElements{
public:
    __aicore__ inline KernelScatterElements() {}
    __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, uint32_t lastdim, uint32_t totalLength, uint32_t ALIGN_NUM, uint32_t block_size, uint32_t core_size, uint32_t core_remain, uint32_t var_shape, uint32_t var_bats, uint32_t var_tier, uint32_t var_rows, uint32_t var_cols, uint32_t dim_shape, uint32_t loop_h, uint32_t loop_i, uint32_t loop_j, uint32_t loop_k, int axis, uint32_t strcase
    ,TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->blockLength = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
        this->tileLength = block_size;
        this->lastdim = lastdim;
        this->totalLength = totalLength;
        this->ALIGN_NUM = ALIGN_NUM;
        this->axis = axis;
        this->strcase = strcase;
        this->var_shape = var_shape;
        this->var_bats = var_bats;
        this->var_tier = var_tier;
        this->var_rows = var_rows;
        this->var_cols = var_cols;
        this->dim_shape = dim_shape;
        this->loop_h = loop_h;
        this->loop_i = loop_i;
        this->loop_j = loop_j;
        this->loop_k = loop_k;
        // this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);

        auto startPointer = core_size * GetBlockIdx();
        auto bufferlength = this->blockLength;
        // if(this->dim_shape == 1){
                // get start index for current core, core parallel
        xGm.SetGlobalBuffer((__gm__ T_VAR*)var, totalLength);//x为初始self,即初始的数组或矩阵
        yGm.SetGlobalBuffer((__gm__ T_INDICES*)indices + startPointer, bufferlength); //y为变更的索引index
        dGm.SetGlobalBuffer((__gm__ T_VAR*)updates + startPointer * this->lastdim, bufferlength * this->lastdim); //d为变更的内容，即更新的值

        this->tileNum = this->lastdim / this->tileLength + (this->lastdim % this->tileLength > 0);
        pipe =pipeIn;
        // pipe alloc memory to queue, the unit is Bytes
        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(T_VAR));
        // pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(T_INDICES));
        pipe->InitBuffer(inQueueD, BUFFER_NUM, this->tileLength * sizeof(T_VAR));
        pipe->InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(T_VAR));   //z为根据变更内容，以索引对初始数组，矩阵进行修改后的最终结果
        if constexpr (std::is_same_v<T_VAR, uint8_t>) {
            pipe->InitBuffer(tmpBuffer1, this->tileLength * sizeof(half));
            pipe->InitBuffer(tmpBuffer2, this->tileLength * sizeof(half));
            pipe->InitBuffer(tmpBuffer3, this->tileLength * sizeof(half));
        }      
    }
    __aicore__ inline void Process() {
        if(this->dim_shape == 1){
            for(int32_t i = 0; i < this->blockLength; ++i){
                int32_t p = yGm.GetValue(i);        //获取原indices的所有的索引值，以p储存
                // tiling strategy, pipeline parallel
                int32_t loopCount = this->tileNum;
                for (int32_t j = 0; j < loopCount-1; ++j) {
                    CopyIn(p, i, j, 0, this->tileLength, 0);
                    Compute(p, i, j, 0, this->tileLength);
                    CopyOut(p, i, j, 0, this->tileLength);
                }
            auto padding = (this->lastdim % this->ALIGN_NUM ? this->ALIGN_NUM - this->lastdim % this->ALIGN_NUM : 0);
            auto length = this->lastdim + padding;
            length = length - this->tileLength * (loopCount - 1);
            CopyIn(p, i, loopCount - 1, 0, length, padding);
            Compute(p, i, loopCount - 1, 0, length);
            CopyOut(p, i, loopCount - 1, 0, length);
            }
        }
        else if(this->dim_shape == 2){
            for(int32_t i = 0; i < loop_i; ++i){
                for(int32_t j = 0; j < loop_j; ++j){
                    int32_t w = yGm.GetValue(i * loop_j + j);
                    int32_t p;
                    if(axis == 0){
                        p = w * var_cols + j;
                    }
                    else if(axis == 1){
                        p = i * var_cols + w;
                    }
                    // int32_t p = w * var_cols + j;
                    int32_t q = i * loop_j + j;
                    int32_t loopCount = this->tileNum;
                    // int32_t tileLength = this->tileLength;
                    // int32_t lastdim = this->lastdim;
                    // int32_t ALIGN_NUM = this->ALIGN_NUM;
                    for (int32_t k = 0; k < loopCount-1; k+=2) {
                        CopyIn(p, q, k, 0, this->tileLength, 0);
                        Compute(p, q, k, 0, this->tileLength);
                        CopyOut(p, q, k, 0, this->tileLength);

                        CopyIn(p, q, k + 1, 0, this->tileLength, 0);
                        Compute(p, q, k + 1, 0, this->tileLength);
                        CopyOut(p, q, k + 1, 0, this->tileLength);
                    }
                    // auto padding = (this->lastdim % this->ALIGN_NUM ? this->ALIGN_NUM - this->lastdim % this->ALIGN_NUM : 0);
                    // auto length = this->lastdim + padding;
                    // length = length - this->tileLength * (loopCount - 1);
                    // CopyIn(p, q, loopCount - 1, 0, length, padding);
                    // Compute(p, q, loopCount - 1, 0, length);
                    // CopyOut(p, q, loopCount - 1, 0, length);
                    if (loopCount % 2 == 1) {
                        int32_t k = loopCount - 1;
                        auto padding = (this->lastdim % this->ALIGN_NUM ? this->ALIGN_NUM - this->lastdim % this->ALIGN_NUM : 0);
                        auto length = this->lastdim+ padding;
                        length = length - tileLength * (loopCount - 1);
                        CopyIn(p, q, k, 0, length, padding);
                        Compute(p, q, k, 0, length);
                        CopyOut(p, q, k, 0, length);
            }
                }
            }    
        }
        else if(this->dim_shape == 3){
            for(int32_t i = 0; i < loop_i; ++i){
                for(int32_t j = 0; j < loop_j; ++j){
                    for(int32_t k = 0; k < loop_k; ++k){
                        int32_t w = yGm.GetValue(i * loop_j * loop_k + j * loop_k + k);
                        int32_t p;
                        if(axis == 0){
                            p = w * var_rows * var_cols + j * var_cols + k;
                        }
                        else if(axis == 1){
                            p = i * var_rows * var_cols + w * var_cols + k;
                        }
                        else if(axis == 2){
                            p = i * var_rows * var_cols + j * var_cols + w;
                        }
                        // int32_t p = w * var_cols + j;
                        int32_t q = i * loop_j * loop_k + j * loop_k + k;
                        int32_t loopCount = this->tileNum;
                        for (int32_t t = 0; t < loopCount-1; ++t) {
                            CopyIn(p, q, t, 0, this->tileLength, 0);
                            Compute(p, q, t, 0, this->tileLength);
                            CopyOut(p, q, t, 0, this->tileLength);
                        }
                        auto padding = (this->lastdim % this->ALIGN_NUM ? this->ALIGN_NUM - this->lastdim % this->ALIGN_NUM : 0);
                        auto length = this->lastdim + padding;
                        length = length - this->tileLength * (loopCount - 1);
                        CopyIn(p, q, loopCount - 1, 0, length, padding);
                        Compute(p, q, loopCount - 1, 0, length);
                        CopyOut(p, q, loopCount - 1, 0, length);
                    }
                }
            }    
        }
        else if(this->dim_shape == 4){
             for(int32_t h = 0; h < loop_h; ++h){
                for(int32_t i = 0; i < loop_i; ++i){
                    for(int32_t j = 0; j < loop_j; ++j){
                        for(int32_t k = 0; k < loop_k; ++k){
                            int32_t w = yGm.GetValue(h * loop_i * loop_j * loop_k + i * loop_j * loop_k + j * loop_k + k);
                            int32_t p;
                            if(axis == 0){
                                p = w * var_tier * var_rows * var_cols + i * var_rows * var_cols + j * var_cols + k;
                            }
                            else if(axis == 1){
                                p = h * var_tier * var_rows * var_cols + w * var_rows * var_cols + j * var_cols + k;
                            }
                            else if(axis == 2){
                                p = h * var_tier * var_rows * var_cols + i * var_rows * var_cols + w * var_cols + k;
                            }
                            else{
                                p = h * var_tier * var_rows * var_cols + i * var_rows * var_cols + j * var_cols + w;
                            }
                            // int32_t p = w * var_cols + j;
                            int32_t q = h * loop_i * loop_j * loop_k + i * loop_j * loop_k + j * loop_k + k;
                            int32_t loopCount = this->tileNum;
                            for (int32_t t = 0; t < loopCount-1; ++t) {
                                CopyIn(p, q, t, 0, this->tileLength, 0);
                                Compute(p, q, t, 0, this->tileLength);
                                CopyOut(p, q, t, 0, this->tileLength);
                            }
                            auto padding = (this->lastdim % this->ALIGN_NUM ? this->ALIGN_NUM - this->lastdim % this->ALIGN_NUM : 0);
                            auto length = this->lastdim + padding;
                            length = length - this->tileLength * (loopCount - 1);
                            CopyIn(p, q, loopCount - 1, 0, length, padding);
                            Compute(p, q, loopCount - 1, 0, length);
                            CopyOut(p, q, loopCount - 1, 0, length);
                        }
                    }
                }
            }
        }
    }
private:
    __aicore__ inline void CopyIn(int32_t p, int32_t i, int32_t j, int32_t k, uint32_t length, uint32_t padding){
        // alloc tensor from queue memory
        LocalTensor<T_VAR> xLocal = inQueueX.AllocTensor<T_VAR>();  //初始var
        LocalTensor<T_VAR> dLocal = inQueueD.AllocTensor<T_VAR>();  //加入变更值，src,value,scalar
        DataCopy(xLocal, xGm[p * this->lastdim + j * this->tileLength], length);
        DataCopy(dLocal, dGm[i * this->lastdim + j * this->tileLength], length);
        if(this->strcase == 1){

            //Duplicate(dLocal,(T_VAR)(0.0),paddings);
            for(int m = length-padding;m < length;m++){
                dLocal.SetValue(m, 0);     
            }

        }
        else if(this->strcase == 2){
            for(int m = length-padding;m < length; m++){
                dLocal.SetValue(m, 0);     
            }

            if  constexpr(std::is_same_v<T_VAR,half>|std::is_same_v<T_VAR,float>
            | std::is_same_v<T_VAR,int32_t> ){

            
            // for(int n = 0; n < length-padding; n++){
            //     xLocal.SetValue(n, 0);
               T_VAR zero = 0.0f;
               Duplicate(xLocal,zero,length-padding);
            }

            else{
                 for(int n = 0; n < length-padding; n++){
                    
                    xLocal.SetValue(n, 0);

            }

            }
        }             
        // enque input tensors to VECIN queue
        inQueueD.EnQue(dLocal);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t p, int32_t i, int32_t j, int32_t k, uint32_t length)
    {
       // deque input tensors from VECIN queue
        LocalTensor<T_VAR> xLocal = inQueueX.DeQue<T_VAR>();
        LocalTensor<T_VAR> dLocal = inQueueD.DeQue<T_VAR>();
        LocalTensor<T_VAR> zLocal = outQueueZ.AllocTensor<T_VAR>();

        // Duplicate(self, static_cast<DTYPE_VAR_REF>(0), rowW);
        if constexpr (std::is_same_v<T_VAR, uint8_t>) {
            LocalTensor<half> dh = tmpBuffer1.Get<half>();
            LocalTensor<half> zh = tmpBuffer2.Get<half>();
            LocalTensor<half> xh = tmpBuffer3.Get<half>();
            half zero = 0;
            Duplicate(zh, zero, length);
            Cast(xh, xLocal, RoundMode::CAST_NONE, length);
            Cast(dh, dLocal, RoundMode::CAST_NONE, length);
            Add(xh, xh, dh, length);    //原数组，矩阵根据索引加上变更的新值，
            Add(zh, xh, zh ,length);    //加上初始化后的0矩阵
            Cast(zLocal, zh, RoundMode::CAST_NONE, length);
        }
        else{
            // T_VAR zero = 0;
            // Duplicate(zLocal, zero, length);
            Add(zLocal, xLocal, dLocal, length);
            // Add(zLocal, xLocal, zLocal, length);
        }
        outQueueZ.EnQue<T_VAR>(zLocal);
        // inQueueX.EnQue<T_VAR>(xLocal);
        // free input tensors for reuse
        inQueueX.FreeTensor(xLocal);
        inQueueD.FreeTensor(dLocal); 
    }
    __aicore__ inline void CopyOut(int32_t p, int32_t i, int32_t j, int32_t k, uint32_t length)
    {
        // deque output tensor from VECOUT queue
        LocalTensor<T_VAR> zLocal = outQueueZ.DeQue<T_VAR>();
        // copy progress_th tile from local tensor to global tensor
      
        DataCopy(xGm[p * this->lastdim + j * this->tileLength], zLocal, length);
        outQueueZ.FreeTensor(zLocal);
    }
private:
    TPipe *pipe;
    TBuf<QuePosition::VECCALC> tmpBuffer1, tmpBuffer2, tmpBuffer3;
    // create queues for input, in this case depth is equal to buffer num
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueD;
    // create queue for output, in this case depth is equal to buffer num
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
    GlobalTensor<T_VAR> xGm;
    GlobalTensor<T_INDICES> yGm;
    GlobalTensor<T_VAR> dGm;
    uint32_t lastdim;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t totalLength;
    uint32_t ALIGN_NUM;
    uint32_t var_shape;
    uint32_t var_bats;
    uint32_t var_tier;
    uint32_t var_rows;
    uint32_t var_cols;
    uint32_t dim_shape;
    uint32_t loop_h;
    uint32_t loop_i;
    uint32_t loop_j;
    uint32_t loop_k;
    int axis;
    uint32_t strcase;
};
extern "C" __global__ __aicore__ void scatter_elements(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    // if(TILING_KEY_IS(1)){
    //     KernelScatterElements1<DTYPE_VAR, DTYPE_INDICES, DTYPE_UPDATES> op;
    //     op.Init(var, indices, updates, tiling_data.lastdim, tiling_data.totalLength, tiling_data.ALIGN_NUM, tiling_data.block_size, tiling_data.core_size, tiling_data.core_remain, tiling_data.dim_shape, tiling_data.loop_i, tiling_data.loop_j, tiling_data.loop_k, tiling_data.axis, tiling_data.strcase);
    //     op.Process();
    // }
    // else if(TILING_KEY_IS(2)){

    // }
    // else if(TILING_KEY_IS(3)){

    // }
    TPipe pipe;
    KernelScatterElements<DTYPE_VAR, DTYPE_INDICES, DTYPE_UPDATES> op;
    op.Init(var, indices, updates, tiling_data.lastdim, tiling_data.totalLength, tiling_data.ALIGN_NUM, tiling_data.block_size, tiling_data.core_size, tiling_data.core_remain, 
    tiling_data.var_shape, tiling_data.var_bats, tiling_data.var_tier, tiling_data.var_rows, tiling_data.var_cols, tiling_data.dim_shape, tiling_data.loop_h, tiling_data.loop_i, tiling_data.loop_j, tiling_data.loop_k, tiling_data.axis, tiling_data.strcase,&pipe);
    op.Process();
}